# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: lucasyoo

import logging
import git
import os
import sys
import time
import collections
from typing import List, Dict, Optional
from absl import app, flags
import ujson
import datetime
from scipy.stats import norm
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import cc.appcoin2.strategy.lm.kline_generator2 as klg2
from coin.base.datetime_util import (to_datetime, to_timestamp_int)

sys.path.append(os.path.expandvars('$COIN_DEPLOY_REPO'))
os.environ['keyword'] = 'sfty_delta'  # just for hack. No meaning
import python.coin_deploy.vmm.fabfile as vmm_fabfile

matplotlib.use("Agg")
os.environ["TZ"] = "GMT"
time.tzset()

FLAGS = flags.FLAGS

Data = collections.namedtuple("Data", ["ts", "product_str", "price", "reserve", "hypo_reserve"])

os.environ["TZ"] = "GMT"
time.tzset()

VMM_CONFIG_DIRECTORY = os.path.expandvars("$COIN_REPO/python/coin_deploy/vmm/config")
MONITOR_CONFIG_DIRECTORY = os.path.expandvars(
    "$COIN_REPO/python/coin_deploy/support_day1mm/team_flow_monitor")


def get_option_pnl(start_time: datetime.datetime, end_time: datetime.datetime, strategy_name: str,
                   base: str):
  product_strs = {
      data.product_str for data in query_coin_deploy_executor_reserve_from_git(strategy_name, base)
  }
  final_df = pd.DataFrame()
  for product_str in product_strs:
    kline_df = query_interval_feed(product_str, start_time, end_time, resolution='1d')
    final_df = pd.concat([final_df, kline_df])
  if final_df.empty:
    return pd.DataFrame(), pd.DataFrame(), {}
  final_df['ts'] = pd.to_datetime(final_df['klineTimestamp'])
  mean_df = final_df.groupby(['ts'])[['open', 'close']].mean()
  mean_df['base'] = base
  mean_df.reset_index(inplace=True, names="ts")
  strikes = []
  maturities = []
  volatilities = []
  rfs = []
  pos_offset = {}
  exec_reserve = {}
  with open(
      os.path.join(os.path.expandvars('$COIN_DEPLOY_REPO'),
                   f"vmm/config/{strategy_name}.json")) as f:
    strat_config = ujson.load(f)
    for product_str, config in strat_config["oc_config"]["configs"].items():
      for k, v in config["contracts"].items():
        if k == "owner":
          continue
        strikes.append(float(v['strike_price_num']) / float(v['strike_price_denom']))
        maturities.append(datetime.datetime.strptime(v["expiry"], "%Y-%m-%d %H:%M:%S"))
      rfs.append(config["risk_free_rate"])
      volatilities.append(config['volatility'])
      pos_offset[f"{strategy_name}|{product_str.split(':')[-1].split('-')[0]}"] = config[
          'pos_offset']
      exec_reserve[f"{strategy_name}|{product_str.split(':')[-1].split('-')[0]}"] = strat_config[
          "executors"][product_str]["reserve"]
  final_reserve = {}
  for k, v in pos_offset.items():
    final_reserve[k] = v + exec_reserve[k]
  assert len(set(strikes)) == 1 and len(set(maturities)) == 1 and len(set(rfs)) == 1 and len(
      set(volatilities)) == 1
  mean_df['strike_price'] = strikes[0]
  mean_df['maturity'] = maturities[0]
  mean_df['rf'] = rfs[0]
  mean_df['volatility'] = volatilities[0]
  mean_df['call_price_diff'] = mean_df.apply(call_option_price_diff, axis=1)
  return mean_df, kline_df, final_reserve


def query_option_num_info():
  with open(os.path.join(MONITOR_CONFIG_DIRECTORY, 'option_num.json')) as f:
    config = ujson.load(f)
  return config


def call_option_price_diff(row):
  close = row['close']
  open = row['open']
  open_ts = row['ts']
  close_ts = row['ts'] + datetime.timedelta(days=1)
  rf = row['rf']
  maturity = row['maturity']
  strike_price = row['strike_price']
  volatility = row['volatility']
  call_price_diff = call_option_price(close, close_ts, maturity, volatility, strike_price, rf) - \
              call_option_price(open, open_ts, maturity, volatility, strike_price, rf)
  return call_price_diff


def call_option_price(spot_price: float, dt: float, maturity: datetime.datetime,
                      implied_volatility: float, strike_price: float,
                      risk_free_rate: float) -> float:
  """
    Calculate the price of a call option using the Black-Scholes model.
    :param implied_volatility: Implied volatility of option (abs value, not percent unit). e.g. 1.0
    :param strike_price: The strike price of the option.
    :param maturity: The maturity of the option
    :param spot_price: The spot price of the underlying asset.
    :param risk_free_rate: The risk free rate (absolute value, not percent unit). e.g. 0.01
    :param call_num: The number of calls.
    :return: The price of the call option.
    """
  assert dt <= maturity, f'dt: {dt} > maturity: {maturity}'
  ttm = (maturity - dt).total_seconds() / (365 * 24 * 60 * 60)
  d1 = (np.log(spot_price / strike_price) +
        (risk_free_rate + 0.5 * implied_volatility**2) * ttm) / (implied_volatility * np.sqrt(ttm))
  d2 = d1 - implied_volatility * np.sqrt(ttm)
  call_price = spot_price * norm.cdf(d1) - strike_price * np.exp(
      -risk_free_rate * ttm) * norm.cdf(d2)
  return call_price


def query_coin_deploy_executor_reserve_from_git(strategy: str, base: str):
  repo = git.Repo(os.path.expandvars('$COIN_DEPLOY_REPO'))
  commits = list(repo.iter_commits(paths=f'vmm/config/{strategy}.json'))
  commits.reverse()
  executor_reserve_list: List[Optional[Data]] = []
  last_exec_reserve_dict = {}
  for commit in commits:
    commit_time = commit.committed_datetime
    utc_commit_time = commit_time.astimezone(datetime.timezone.utc).replace(tzinfo=None)
    tree = commit.tree
    for blob in tree.traverse():
      if blob.path == f'vmm/config/{strategy}.json':
        try:
          strat_config = ujson.loads(blob.data_stream.read().decode('utf-8'))
        except ujson.JSONDecodeError:
          logging.error(f'[{blob.path}] {utc_commit_time} failed to read json'
                        f'{commit.hexsha} {commit.message}')
          continue
        for product_str, exec_config in strat_config['executors'].items():
          if base in product_str:
            exec_reserve = exec_config['reserve']
            if exec_reserve != last_exec_reserve_dict.get(product_str):
              executor_reserve_list.append(
                  Data(utc_commit_time, product_str, np.nan, np.nan, exec_reserve))
              logging.info(f'[{product_str}] {utc_commit_time} hypo reserve changed from '
                           f'{last_exec_reserve_dict.get(product_str, 0.0)} to {exec_reserve}')
              last_exec_reserve_dict[product_str] = exec_reserve
          else:
            logging.info(f'[{product_str}] {utc_commit_time} not base')
  return executor_reserve_list


def get_mea(product_str: str):
  market, exchange, product = product_str.split(':', 2)
  me = f'{market}.{exchange}'
  if me == 'Spot.Binance':
    return me + '.v1'

  if me == 'Futures.Binance':
    if 'USDT' in product:
      return me + '.v1'
    else:
      return me + '.v1-delivery'

  if me == 'Spot.Ftx':
    return me + '.v1'

  if me == 'Futures.Ftx':
    return me + '.v1'

  if me == "Spot.Huobi":
    return me + '.v1'

  if me == "Spot.Gateio":
    return me + '.v4'

  if me == "Spot.Mexc":
    return me + '.v3'

  if me == "Spot.Okex":
    return me + '.v5'

  if me == "Spot.Kucoin":
    return me + '.v1'

  if me == "Spot.Coinone":
    return me + '.v2'

  if me == "Spot.Mercado":
    return me + '.v4'

  if me == "Spot.Bitmart":
    return me + '.v3'
  if me == "Spot.Bithumb":
    return me + '.v2'
  if me == "Spot.Bitbank":
    return me + '.v1'
  if me == "Spot.Bybit":
    return me + '.v3'

  if me == "Spot.Uniswap":
    pi_v2_dir = os.path.expandvars("${COIN_REPO}/data/coin2/product_info/Spot.Uniswap.v2.json")
    pi_v3_dir = os.path.expandvars("${COIN_REPO}/data/coin2/product_info/Spot.Uniswap.v3.json")
    with open(pi_v3_dir) as f:
      pi_v3 = pd.DataFrame(ujson.load(f)['product_infos'])
    with open(pi_v2_dir) as f:
      pi_v2 = pd.DataFrame(ujson.load(f)['product_infos'])
    if product in pi_v3['symbol'].tolist():  # use Uniswap V3 if available
      return me + '.v3'
    elif product in pi_v2['symbol'].tolist():
      return me + '.v2'
  if me == "Spot.Phemex":
    return me + '.v1'


def query_interval_feed(product_str: str,
                        start_dt: datetime.datetime,
                        end_dt: datetime.datetime,
                        resolution: str = '1m'):
  final_df = pd.DataFrame()
  mea = get_mea(product_str)
  _, _, symbol = product_str.split(":")
  if mea == "Spot.Mexc.v3":
    if start_dt < datetime.datetime(2023, 2, 24) and end_dt >= datetime.datetime(2023, 2, 24):
      tmp_mea = mea.replace("3", "2")
      tmp_df = klg2.get_kline_from_interval_feed(start_dt,
                                                 datetime.datetime(2023, 2, 24),
                                                 tmp_mea, [symbol],
                                                 return_df=True,
                                                 resolution=resolution)[0].get(symbol)
      if tmp_df is None:
        tmp_df = pd.DataFrame()
      else:
        tmp_df['product_str'] = product_str
      final_df = pd.concat([final_df, tmp_df])
      start_dt = datetime.datetime(2023, 2, 24)
    elif start_dt < datetime.datetime(2023, 2, 24) and end_dt < datetime.datetime(2023, 2, 24):
      mea = mea.replace("3", "2")
  kline_df = klg2.get_kline_from_interval_feed(start_dt,
                                               end_dt,
                                               mea, [symbol],
                                               return_df=True,
                                               resolution=resolution)[0].get(symbol)

  final_df = pd.concat([final_df, kline_df]) if kline_df is not None else final_df
  if final_df.empty:
    return final_df
  final_df['product_str'] = product_str
  final_df['close'].ffill(inplace=True)
  return final_df


class PnlAccounting(object):
  def __init__(self, symbol: str, start_dt: datetime.datetime, end_dt: datetime.datetime):
    assert end_dt - start_dt > datetime.timedelta(
        hours=24), 'start_dt and end_dt must be at least 24 hours apart.'
    assert start_dt.time() == datetime.time(0, 0, 0) and end_dt.time() == datetime.time(
        0, 0, 0), 'start_dt and end_dt must be at midnight.'
    self.symbol = symbol
    self.start_dt = start_dt
    self.end_dt = end_dt
    self.start_dt_int = to_timestamp_int(start_dt)
    self.end_dt_int = to_timestamp_int(end_dt)
    self.query_corp_info()
    self.query_klines()

  @staticmethod
  def call_option_price(implied_volatility: float, strike_price: float, maturity: datetime.datetime,
                        dt: datetime.datetime, spot_price: np.array,
                        risk_free_rate: float) -> np.array:
    """
    Calculate the price of a call option using the Black-Scholes model.
    :param implied_volatility: Implied volatility of option (abs value, not percent unit). e.g. 1.0
    :param strike_price: The strike price of the option.
    :param maturity: The maturity of the option
    :param spot_price: The spot price of the underlying asset.
    :param risk_free_rate: The risk free rate (absolute value, not percent unit). e.g. 0.01
    :param call_num: The number of calls.
    :return: The price of the call option.
    """
    assert dt <= maturity, f'dt: {dt} > maturity: {maturity}'
    ttm = (maturity - dt).total_seconds() / (365 * 24 * 60 * 60)
    d1 = (np.log(spot_price / strike_price) +
          (risk_free_rate + 0.5 * implied_volatility**2) * ttm) / (implied_volatility *
                                                                   np.sqrt(ttm))
    d2 = d1 - implied_volatility * np.sqrt(ttm)
    call_price = spot_price * norm.cdf(d1) - strike_price * np.exp(
        -risk_free_rate * ttm) * norm.cdf(d2)
    return call_price

  def query_corp_info(self):
    with open(os.path.join(MONITOR_CONFIG_DIRECTORY, 'option_num.json')) as f:
      config = ujson.load(f)
    self.corp_option_num = config[self.symbol]["total"]

  def query_klines(self):
    result = pd.DataFrame()
    for k, _ in self.exchange_call_info.items():
      tmp_df = query_interval_feed(k, self.start_dt, self.end_dt)
      result = pd.concat([result, tmp_df])
    self.kline_df = result

  def calc_corp_option_pnl(self) -> pd.Series:
    """
    :return: pd.Series of pnl
    """
    # 1, Get Average Price of Spot
    avg_close_price = []
    avg_open_price = []
    for _, kline_df in self.kline_df.items():
      tmp_df = kline_df.loc[(kline_df['klineTimestamp'] >= self.start_dt_int) &
                            (kline_df['klineTimestamp'] <= self.end_dt_int),
                            ['klineTimestamp', 'close', 'open']].copy()
      tmp_df.set_index('klineTimestamp', inplace=True)
      avg_close_price.append(tmp_df['close'])
      avg_open_price.append(tmp_df['open'])

    avg_close_df = pd.concat(avg_close_price, axis=1).mean(axis=1)
    avg_open_df = pd.concat(avg_open_price, axis=1).mean(axis=1)

    # 2. Get option pricing config from coin_deploy
    assert self.strat_option_num, 'strat_option_num is empty'
    im_vol = None
    rf = None
    strike = None
    maturity = None
    for _, v in self.strat_option_num:
      im_vol = v["volatility"]
      rf = v["risk_free_rate"]
      strike = (v["strike_price_num"] if v.get("report_strike_price_denom") is None else
                v["report_strike_price_denom"]) / (v["strike_price_denum"]
                                                   if v.get("report_strike_price_denom") is None
                                                   else v["report_strike_price_denom"])

    # close_call_price = self.


def main(_):
  start_time = datetime.datetime.strptime(FLAGS.start_time, '%Y%m%d')
  end_time = datetime.datetime.strptime(FLAGS.end_time, '%Y%m%d')


if __name__ == "__main__":
  flags.DEFINE_string('start_time', '20230101', '%Y%m%d')
  flags.DEFINE_string('end_time', '20230331', '%Y%m%d')
  flags.DEFINE_string('strategy_name', 'ALL',
                      'strategy name e.g.vmm_ccx_delta_hedge, if ALL, run every vmm_*_delta_hedge')
  flags.DEFINE_string('output_dir', 'sim_result/accounting', 'output dir')
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)